#!/usr/bin/env bash
cmd=$( basename "$0" )
img="rknn-toolkit2:2.3.0-cp38"
dockerName="ljs-rknn-docker"

gpu=false
function usage() {
    echo "Usage: ${cmd} [-i | -n | -g | -h]" 
    echo "      启动docker"
    echo "      -i: <docker_img>:<docker_tag>, 缺省: ${img}"
    echo "      -n: docker容器名称，缺省: ${dockerName}"
    echo "      -g: 增加gpu支持"
    echo "      -h: 输出此信息"
    exit 1
}

shiftCnt=0
while getopts "i:n:gh" arg
do
    case "${arg}" in 
        i)
            img="${OPTARG}"
            shiftCnt=$(( ${shiftCnt} + 2 ))
            ;;
        n)
            dockerName="${OPTARG}"
            shiftCnt=$(( ${shiftCnt} + 2 ))
            ;;
        g)
            gpu=true
            shiftCnt=$(( ${shiftCnt} + 1 ))
            ;;
        *)
            usage
            ;;
    esac
done
shift ${shiftCnt}

if [ $# -gt 0 ]
then
    usage
fi

echo "Run image ${img}..."

hostDir="${HOME}/docker_homes/rknn_home"
etcDir="${HOME}/sys/etc_rknn"
user="1172"
group="502"
clientHome="/home/rknn"
workDir="${HOME}/work/polarbird-rknn"

echo run docker image ${img}
# 先删除旧的container
docker container rm -f ${dockerName} 2> /dev/null

if ${gpu}
then
    gpuOption='--gpus all' 
else
    gpuOption=''
fi

docker run \
    -u ${user}:${group} \
    ${gpuOption} \
    --net host \
    -e DISPLAY=${DISPLAY} \
    --rm \
    -v /tmp/.X11-unix/:/tmp/.X11-unix/ \
    -v ${HOME}/.Xauthority:${clientHome}/.Xauthority \
    -v /dev/bus/usb:/dev/bus/usb \
    -v /usr/local/cuda-12.5:/usr/local/cuda-12.5 \
    -v ${etcDir}/group:/etc/group:ro \
    -v ${etcDir}/passwd:/etc/passwd:ro \
    -v ${etcDir}/shadow:/etc/shadow:ro \
    -v ${etcDir}/sudoers:/etc/sudoers \
    -v ${hostDir}:${clientHome} \
    -v ${workDir}/rknn_model_zoo-v2.3.0-2024-11-08:${clientHome}/rknn_model_zoo \
    -v ${workDir}/rknn-toolkit2-v2.3.0-2024-11-08:${clientHome}/rknn-tookkit2 \
    -v ${HOME}/bin:${clientHome}/bin \
    -v ${HOME}/exer:${clientHome}/exer \
    -v ${HOME}/download:${clientHome}/download \
    -v ${HOME}/tmp:${clientHome}/tmp \
    -v ${HOME}/datasets:${clientHome}/datasets \
    -w ${clientHome} \
    --name "${dockerName}" \
    -it \
    ${img}
